package com.gitee.xhs.operatorState.listState;

import com.gitee.xhs.sampleRecod.transactionRecord.Transaction;
import com.gitee.xhs.sampleRecod.transactionRecord.TransactionIterator;
import com.gitee.xhs.sampleRecod.transactionRecord.jmockTool.TransactionLoopOneToFiveMockerInterceptor;
import com.gitee.xhs.sourceIterator.RateLimitedIterator;
import com.gitee.xhs.util.StringTemplateUtil;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.sql.Timestamp;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * @author BigHead
 * @date 2022/7/1 14:38
 * @description OperatorListStateSource
 * <p>
 * 数据生成时间从1990年开始
 * 每增加一个并发则增加一年用作测试
 */
public class OperatorListStateSource extends RichParallelSourceFunction<Transaction> implements CheckpointedFunction {

    private static final String START_TIMESTAMP = "${YEAR}-01-01 00:00:00";
    private static final Integer START_YEAR = 1990;
    private static final String YEAR_KEY = "YEAR";

    private Iterator<Transaction> transactionIterator;

    private volatile boolean isRunning = true;
    private volatile Long recordTimeStamp = 0L;
    private transient ListState<Long> listState;

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        listState.clear();
        listState.add(recordTimeStamp);
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        int parallelismIndex = getRuntimeContext().getIndexOfThisSubtask();
        recordTimeStamp = Timestamp.valueOf(yearUtil(parallelismIndex)).getTime();

        listState = context.getOperatorStateStore().getListState(OperatorListStateConstant.OPERATOR_LIST_STATE_DESCRIPTOR);
        if (context.isRestored()) {
            int stateIndex = 0;
            for (Long stateLong : listState.get()) {
                if (stateIndex == parallelismIndex) {
                    recordTimeStamp = stateLong;
                }
                stateIndex += 1;
            }
        }

        TransactionIterator initTransactionIterator = new TransactionIterator(
                true,
                new TransactionLoopOneToFiveMockerInterceptor(),
                null,
                recordTimeStamp,
                null,
                null
        );

        transactionIterator = new RateLimitedIterator<>(
                initTransactionIterator
        );

    }

    @Override
    public void run(SourceContext<Transaction> ctx) throws Exception {
        int parallelismIndex = getRuntimeContext().getIndexOfThisSubtask();
        while (isRunning && transactionIterator.hasNext()) {
            Transaction transaction = transactionIterator.next();
            transaction.setTaskIndexNum(parallelismIndex);
            recordTimeStamp = transaction.getTimestamp();
            ctx.collect(transaction);
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
    }

    private String yearUtil(int parallelismIndex) {
        Map<String, String> templateMap = new HashMap<>(1);
        templateMap.put(
                YEAR_KEY,
                START_YEAR + parallelismIndex + ""
        );
        return StringTemplateUtil.render(
                START_TIMESTAMP,
                templateMap
        );
    }
}
